package main

import (
	"fmt"
	"image"
	"image/color"
	"image/jpeg"
	"log"
	"os"
)

var color_map_r = [...]float32{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.00588235294117645, 0.02156862745098032, 0.03725490196078418, 0.05294117647058827, 0.06862745098039214, 0.084313725490196, 0.1000000000000001, 0.115686274509804, 0.1313725490196078, 0.1470588235294117, 0.1627450980392156, 0.1784313725490196, 0.1941176470588235, 0.2098039215686274, 0.2254901960784315, 0.2411764705882353, 0.2568627450980392, 0.2725490196078431, 0.2882352941176469, 0.303921568627451, 0.3196078431372549, 0.3352941176470587, 0.3509803921568628, 0.3666666666666667, 0.3823529411764706, 0.3980392156862744, 0.4137254901960783, 0.4294117647058824, 0.4450980392156862, 0.4607843137254901, 0.4764705882352942, 0.4921568627450981, 0.5078431372549019, 0.5235294117647058, 0.5392156862745097, 0.5549019607843135, 0.5705882352941174, 0.5862745098039217, 0.6019607843137256, 0.6176470588235294, 0.6333333333333333, 0.6490196078431372, 0.664705882352941, 0.6803921568627449, 0.6960784313725492, 0.7117647058823531, 0.7274509803921569, 0.7431372549019608, 0.7588235294117647, 0.7745098039215685, 0.7901960784313724, 0.8058823529411763, 0.8215686274509801, 0.8372549019607844, 0.8529411764705883, 0.8686274509803922, 0.884313725490196, 0.8999999999999999, 0.9156862745098038, 0.9313725490196076, 0.947058823529412, 0.9627450980392158, 0.9784313725490197, 0.9941176470588236, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0.9862745098039216, 0.9705882352941178, 0.9549019607843139, 0.93921568627451, 0.9235294117647062, 0.9078431372549018, 0.892156862745098, 0.8764705882352941, 0.8607843137254902, 0.8450980392156864, 0.8294117647058825, 0.8137254901960786, 0.7980392156862743, 0.7823529411764705, 0.7666666666666666, 0.7509803921568627, 0.7352941176470589, 0.719607843137255, 0.7039215686274511, 0.6882352941176473, 0.6725490196078434, 0.6568627450980391, 0.6411764705882352, 0.6254901960784314, 0.6098039215686275, 0.5941176470588236, 0.5784313725490198, 0.5627450980392159, 0.5470588235294116, 0.5313725490196077, 0.5156862745098039, 0.5}
var color_map_g = [...]float32{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.001960784313725483, 0.01764705882352935, 0.03333333333333333, 0.0490196078431373, 0.06470588235294117, 0.08039215686274503, 0.09607843137254901, 0.111764705882353, 0.1274509803921569, 0.1431372549019607, 0.1588235294117647, 0.1745098039215687, 0.1901960784313725, 0.2058823529411764, 0.2215686274509804, 0.2372549019607844, 0.2529411764705882, 0.2686274509803921, 0.2843137254901961, 0.3, 0.3156862745098039, 0.3313725490196078, 0.3470588235294118, 0.3627450980392157, 0.3784313725490196, 0.3941176470588235, 0.4098039215686274, 0.4254901960784314, 0.4411764705882353, 0.4568627450980391, 0.4725490196078431, 0.4882352941176471, 0.503921568627451, 0.5196078431372548, 0.5352941176470587, 0.5509803921568628, 0.5666666666666667, 0.5823529411764705, 0.5980392156862746, 0.6137254901960785, 0.6294117647058823, 0.6450980392156862, 0.6607843137254901, 0.6764705882352942, 0.692156862745098, 0.7078431372549019, 0.723529411764706, 0.7392156862745098, 0.7549019607843137, 0.7705882352941176, 0.7862745098039214, 0.8019607843137255, 0.8176470588235294, 0.8333333333333333, 0.8490196078431373, 0.8647058823529412, 0.8803921568627451, 0.8960784313725489, 0.9117647058823528, 0.9274509803921569, 0.9431372549019608, 0.9588235294117646, 0.9745098039215687, 0.9901960784313726, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0.9901960784313726, 0.9745098039215687, 0.9588235294117649, 0.943137254901961, 0.9274509803921571, 0.9117647058823528, 0.8960784313725489, 0.8803921568627451, 0.8647058823529412, 0.8490196078431373, 0.8333333333333335, 0.8176470588235296, 0.8019607843137253, 0.7862745098039214, 0.7705882352941176, 0.7549019607843137, 0.7392156862745098, 0.723529411764706, 0.7078431372549021, 0.6921568627450982, 0.6764705882352944, 0.6607843137254901, 0.6450980392156862, 0.6294117647058823, 0.6137254901960785, 0.5980392156862746, 0.5823529411764707, 0.5666666666666669, 0.5509803921568626, 0.5352941176470587, 0.5196078431372548, 0.503921568627451, 0.4882352941176471, 0.4725490196078432, 0.4568627450980394, 0.4411764705882355, 0.4254901960784316, 0.4098039215686273, 0.3941176470588235, 0.3784313725490196, 0.3627450980392157, 0.3470588235294119, 0.331372549019608, 0.3156862745098041, 0.2999999999999998, 0.284313725490196, 0.2686274509803921, 0.2529411764705882, 0.2372549019607844, 0.2215686274509805, 0.2058823529411766, 0.1901960784313728, 0.1745098039215689, 0.1588235294117646, 0.1431372549019607, 0.1274509803921569, 0.111764705882353, 0.09607843137254912, 0.08039215686274526, 0.06470588235294139, 0.04901960784313708, 0.03333333333333321, 0.01764705882352935, 0.001960784313725483, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}
var color_map_b = [...]float32{0.5, 0.5156862745098039, 0.5313725490196078, 0.5470588235294118, 0.5627450980392157, 0.5784313725490196, 0.5941176470588235, 0.6098039215686275, 0.6254901960784314, 0.6411764705882352, 0.6568627450980392, 0.6725490196078432, 0.6882352941176471, 0.7039215686274509, 0.7196078431372549, 0.7352941176470589, 0.7509803921568627, 0.7666666666666666, 0.7823529411764706, 0.7980392156862746, 0.8137254901960784, 0.8294117647058823, 0.8450980392156863, 0.8607843137254902, 0.8764705882352941, 0.892156862745098, 0.907843137254902, 0.9235294117647059, 0.9392156862745098, 0.9549019607843137, 0.9705882352941176, 0.9862745098039216, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0.9941176470588236, 0.9784313725490197, 0.9627450980392158, 0.9470588235294117, 0.9313725490196079, 0.915686274509804, 0.8999999999999999, 0.884313725490196, 0.8686274509803922, 0.8529411764705883, 0.8372549019607844, 0.8215686274509804, 0.8058823529411765, 0.7901960784313726, 0.7745098039215685, 0.7588235294117647, 0.7431372549019608, 0.7274509803921569, 0.7117647058823531, 0.696078431372549, 0.6803921568627451, 0.6647058823529413, 0.6490196078431372, 0.6333333333333333, 0.6176470588235294, 0.6019607843137256, 0.5862745098039217, 0.5705882352941176, 0.5549019607843138, 0.5392156862745099, 0.5235294117647058, 0.5078431372549019, 0.4921568627450981, 0.4764705882352942, 0.4607843137254903, 0.4450980392156865, 0.4294117647058826, 0.4137254901960783, 0.3980392156862744, 0.3823529411764706, 0.3666666666666667, 0.3509803921568628, 0.335294117647059, 0.3196078431372551, 0.3039215686274508, 0.2882352941176469, 0.2725490196078431, 0.2568627450980392, 0.2411764705882353, 0.2254901960784315, 0.2098039215686276, 0.1941176470588237, 0.1784313725490199, 0.1627450980392156, 0.1470588235294117, 0.1313725490196078, 0.115686274509804, 0.1000000000000001, 0.08431372549019622, 0.06862745098039236, 0.05294117647058805, 0.03725490196078418, 0.02156862745098032, 0.00588235294117645, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}

const NumDisparity = 128
const INVALID_DISP_VAL = 0
const halfWindow = 11

func openJpg(fileName string) image.Image {
	file, err := os.Open(fileName)
	if err != nil {
		log.Fatal(err)
	}
	// decode jpeg into image.Image
	img, err := jpeg.Decode(file)
	if err != nil {
		log.Fatal(err)
	}
	file.Close()
	return img
}

func cvtColorRGBA2Gray(src image.Image) *image.Gray {
	bounds := src.Bounds()
	gray := image.NewGray(bounds)

	for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
		for x := bounds.Min.X; x < bounds.Max.X; x++ {
			r, g, b, _ := src.At(x, y).RGBA()
			var pixel color.Gray
			pixel.Y = uint8((r*299 + g*587 + b*114) / 1000)
			gray.SetGray(x, y, pixel)
		}
	}
	return gray

}
func cvtColorGray2RGBA_ColorMap(src *image.Gray) *image.RGBA {
	bounds := src.Bounds()
	rgba := image.NewRGBA(bounds)

	for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
		for x := bounds.Min.X; x < bounds.Max.X; x++ {
			pix := src.GrayAt(x, y).Y * 2
			pixel := color.RGBA{uint8(color_map_r[pix] * 255), uint8(color_map_g[pix] * 255), uint8(color_map_b[pix] * 255), 0}
			rgba.SetRGBA(x, y, pixel)
		}

	}
	return rgba
}

func sobel(src *image.Gray) *image.Gray {
	var SOBEL_CAP uint8 = 34
	bounds := src.Bounds()
	dst := image.NewGray(bounds)
	width := bounds.Max.X - bounds.Min.X
	in_image := src.Pix
	for y := bounds.Min.Y + 1; y < bounds.Max.Y-1; y++ {
		for x := bounds.Min.X + 1; x < bounds.Max.X-1; x++ {
			val := ((in_image[width*(y-1)+x+1] - in_image[width*(y-1)+x-1]) +
				(in_image[width*y+x+1]-in_image[width*y+x-1])*2 +
				(in_image[width*(y+1)+x+1] - in_image[width*(y+1)+x-1]))

			if val > SOBEL_CAP {
				val = SOBEL_CAP
			} else if val < -SOBEL_CAP {
				val = -SOBEL_CAP
			}
			dst.SetGray(x, y, color.Gray{val + SOBEL_CAP})
		}
	}
	return dst
}

func left_bm(left_src, right_src *image.Gray) *image.Gray {

	bounds := left_src.Bounds()
	dst := image.NewGray(bounds)
	imgL := left_src.Pix
	imgR := right_src.Pix
	width := bounds.Max.X - bounds.Min.X
	for y := bounds.Min.Y + halfWindow; y < bounds.Max.Y-halfWindow; y++ {
		for x := bounds.Min.X + halfWindow; x < bounds.Max.X-halfWindow; x++ {
			min_costL := 32767

			dst.SetGray(x, y, color.Gray{INVALID_DISP_VAL})
			var d_range int
			if x-halfWindow-NumDisparity > 0 {
				d_range = NumDisparity
			} else {
				d_range = x - halfWindow
			}
			for d := 0; d < d_range; d++ {
				var win_cost = 0
				var diff = 0
				for y1 := y - halfWindow; y1 <= y+halfWindow; y1++ {
					for x1 := x - halfWindow; x1 <= x+halfWindow; x1++ {
						diff = (int)(imgL[y1*width+x1] - imgR[y1*width+x1-d])
						if diff < 0 {
							diff = -diff
						}
						win_cost += diff
					}
				}
				if win_cost < min_costL {
					min_costL = win_cost
					dst.SetGray(x, y, color.Gray{uint8(d)})
				}
			}
		}
	}
	return dst
}
func right_bm(left_src, right_src *image.Gray) *image.Gray {
	bounds := left_src.Bounds()
	dst := image.NewGray(bounds)
	imgL := left_src.Pix
	imgR := right_src.Pix
	width := bounds.Max.X - bounds.Min.X
	for y := bounds.Min.Y + halfWindow; y < bounds.Max.Y-halfWindow; y++ {
		for x := bounds.Min.X + halfWindow; x < bounds.Max.X-halfWindow; x++ {
			min_costL := 32767

			dst.SetGray(x, y, color.Gray{INVALID_DISP_VAL})
			var d_range int
			if width-x-halfWindow-NumDisparity > 0 {
				d_range = NumDisparity
			} else {
				d_range = width - x - halfWindow
			}
			for d := 0; d < d_range; d++ {
				var win_cost = 0
				var diff = 0
				for y1 := y - halfWindow; y1 <= y+halfWindow; y1++ {
					for x1 := x - halfWindow; x1 <= x+halfWindow; x1++ {
						diff = int(imgR[y1*width+x1] - imgL[y1*width+x1+d])
						if diff < 0 {
							diff = -diff
						}
						win_cost += diff
					}
				}
				if win_cost < min_costL {
					min_costL = win_cost
					dst.SetGray(x, y, color.Gray{uint8(d)})
				}
			}
		}
	}
	return dst
}
func left_right_check(left_src, right_src *image.Gray) *image.Gray {
	bounds := left_src.Bounds()
	dst := image.NewGray(bounds)
	width := bounds.Max.X - bounds.Min.X
	disp_left := left_src.Pix
	disp_right := right_src.Pix
	disp := dst.Pix
	for y := bounds.Min.Y + 1; y < bounds.Max.Y-1; y++ {
		for x := bounds.Min.X + 1; x < bounds.Max.X-1; x++ {
			disp[y*width+x] = disp_left[y*width+x]
			var left = disp_left[y*width+x]
			if x-int(left) > 0 {
				right := disp_right[y*width+x-int(left)]
				dispDiff := int(left - right)
				if dispDiff < 0 {
					dispDiff = -dispDiff
				}
				if dispDiff > 1 {
					disp[y*width+x] = 0
				}

			}
		}
	}
	return dst
}
func main() {
	fmt.Println("Open Images")
	left := openJpg("left1.jpg")
	right := openJpg("right1.jpg")

	fmt.Println("Convert to GRAY")
	left_gray := cvtColorRGBA2Gray(left)
	right_gray := cvtColorRGBA2Gray(right)

	fmt.Println("Sobel")
	left_sobel := sobel(left_gray)
	right_sobel := sobel(right_gray)

	fmt.Println("Left BM")
	left_bm := left_bm(left_sobel, right_sobel)

	fmt.Println("Right BM")
	right_bm := right_bm(left_sobel, right_sobel)

	fmt.Println("Left right check")
	bm := left_right_check(left_bm, right_bm)

	fmt.Println("Convert to RGB")
	rgb_out := cvtColorGray2RGBA_ColorMap(bm)
	out, err := os.Create("stereo_bm.jpg")
	if err != nil {
		log.Fatal(err)
	}
	defer out.Close()

	jpeg.Encode(out, rgb_out, nil)
	fmt.Println("Done. The out file is stereo_bm.jpg")
}
